from langchain_core.messages import HumanMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnableWithMessageHistory
from langchain_openai import ChatOpenAI
from langchain_community.chat_message_histories import ChatMessageHistory

model = ChatOpenAI(model="ictrek/qwen7b:32k",
                   openai_api_key="ollama",
                   openai_api_base="http://10.2.4.31:11434/v1/")

# 提示词模版
prompt_template = ChatPromptTemplate.from_messages([
    ('system', '你是一个乐于助人的助手，用{lang}语言尽你所能的回答所有问题。'),
    MessagesPlaceholder(variable_name="msg") # 历史消息
])

# 得到链
chain = prompt_template | model

# 保存聊天记录
store = {}  # 所有用户的聊天记录到保存到store中。key：sessionId, value: 历史聊天记录对象


# 此函数预期接收一个session_id并返回一个消息历史记录对象
def get_session_history(session_id: str):
    if session_id not in store:
        store[session_id] = ChatMessageHistory()
    return store[session_id]


do_message = RunnableWithMessageHistory(
    chain,
    get_session_history, # 回调函数
    input_messages_key="msg" # 每次聊天时候发送msg的健
)

config = {'configurable': { 'session_id': 'abc123' }} # 定义当前会话的session_id
# 第一轮
resp = do_message.invoke(
    {
        "msg": [HumanMessage(content="你好啊！我是阿华")],
        "lang": "中文"
    },
    # 定义sessionId
    config=config
)
print(resp.content)

# 第二轮
resp = do_message.invoke(
    {
        "msg": [HumanMessage(content="请问：我的名字是什么？")],
        "lang": "中文"
    },
    # 定义sessionId
    config=config
)
print(resp.content)

# 第三轮：返回的数据是流式的
resp = do_message.stream(
    {
        "msg": [HumanMessage(content="请给我讲一个笑话")],
        "lang": "中文"
    },
    # 定义sessionId
    config=config
)
for res in resp:
    # 每次响应是一个token
    print(res.content, end="")




